package org.yagnus.stats.samplers.discrete.withoutreplacement;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

/**
 * 
 * @author hc.busy
 * 
 *         This is a most obvious implementation using arralists and O(n)
 *         get/remove methods, but it should be correct.
 * 
 * 
 * 
 * 
 *         For large samples TODO: Check counts after repeated sample from
 *         unequal probabilities--check against brute force computed expected
 *         value. TODO: Check the order in which an item is drawn after a large
 *         number of takeAll by storing the # of draw rank. (compute draw rank
 *         theoretcially and compare)
 * 
 *         TODO: Perhaps write a more optimal version of TakeAll?
 * 
 */
public class ListSampler<T> extends WithoutReplacementSampler<T> {

	public ListSampler(List<T> t, List<Double> weights) {
		super(t, weights);
	}

	public ListSampler(List<T> t) {
		super(t);
	}

	/**
	 * 
	 * @returns next sample or null on error.
	 */
	@Override
	public T take() {
		if (actualSamps.size() == 0 || curTotal == 0) {
			return null;
		}
		T ret = null;
		double myTotal = 0;
		double d = rng().nextDouble() * curTotal;
		Iterator<T> itr = actualSamps.iterator();
		Iterator<Double> apitr = actualP.iterator();
		while (itr.hasNext()) {
			T nextT = itr.next();
			Double nextP = apitr.next();
			myTotal += nextP;

			if (myTotal >= d) {
				// got our sample.
				itr.remove();
				apitr.remove();
				ret = nextT;
				curTotal -= nextP;
				numberSampled++;
				break;
			}
		}
		return ret;
	}

	List<Double> p;
	List<T> samps;
	List<Double> actualP;
	List<T> actualSamps;
	double curTotal;

	@Override
	protected void _init(List<T> t, List<Double> weights) {
		p = new ArrayList<Double>();
		samps = new ArrayList<T>();
		if (t.size() != weights.size()) {
			throw new IllegalArgumentException(
					"weights not same length as sample space.");
		}
		for (int i = 0; i < t.size(); ++i) {
			T tv = t.get(i);
			if (tv == null) {
				throw new IllegalArgumentException(
						"null weights are not suported.");
			}
			this.p.add(weights.get(i));
			this.samps.add(t.get(i));
		}
		_reset();

	}

	protected void _reset() {
		curTotal = 0;
		actualP = new LinkedList<Double>();
		actualSamps = new LinkedList<T>();
		numberSampled = 0;

		for (Double d : p) {
			double v = d.doubleValue();
			curTotal += v;
			actualP.add(v);
		}

		for (T t : samps) {
			actualSamps.add(t);
		}
	}

	@Override
	public void addSample(T t, double w) {
		samps.add(t);
		p.add(w);
		actualSamps.add(t);
		actualP.add(w);
		curTotal += w;
	}

	@Override
	public void removeSample(T t) {

		// Remove from original samples pool.
		for (int i = 0; i < p.size(); ++i) {
			if (samps.get(i).equals(t)) {
				samps.remove(i);
				p.remove(i);
				--i;
			}
		}

		// remove from current remaining sample pool
		for (int i = 0; i < actualP.size(); ++i) {
			if (samps.get(i).equals(t)) {
				samps.remove(i);
				curTotal -= p.remove(i);
				--i; // keep i same since we deleted last one.
			} // now remove from part of the list outside of the available
				// sample.
		}
	}

	public void reset() {
		_reset();
	}

	@Override
	public int remaining() {
		return samps.size() - numberSampled;
	}
}
